# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
import time

from mindspore import context
from mindspore.train import Model
from mindspore.train.serialization import load_checkpoint
from mindspore.nn.metrics import Accuracy, F1
from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
import numpy as np

from src.config import Config
from src.textrcnn import Net
from util import load_data, get_time_dif, get_args


def test_model(network, test_iter):
    start = time.time()
    print("============== Starting Testing ==============")
    # 定义损失函数
    net_loss = SoftmaxCrossEntropyWithLogits()
    model = Model(network, net_loss, metrics={'Accuracy': Accuracy(), 'F1': F1()})
    result = model.eval(test_iter, dataset_sink_mode=False)
    print('Time usage:{}s'.format(get_time_dif(start)))
    print('The accuracy: {}'.format(result.get('Accuracy')))
    print('The F1: {}'.format(np.mean(result.get('F1'))))


def main():
    # 数据集地址和模型地址
    dataset = get_args()

    # 配置文件
    config = Config(dataset)
    context.set_context(mode=context.GRAPH_MODE,
                        device_target=config.device)

    # 获取词汇表和数据
    vocab, train_iter, test_iter = load_data(config)

    # 设置词汇表的大小
    config.n_vocab = len(vocab)

    # 模型读取
    network = Net(config)

    # 加载模型
    load_checkpoint(config.model_path, network)

    # 模型评估
    test_model(network, test_iter)


if __name__ == '__main__':
    main()
